package geppetto.cat.analysis;


import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.alignments.output.AlignerOutputLatex;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.AbstractModel;
import geppetto.cat.models.AgreementHMM;
import geppetto.cat.models.HMM;

import java.io.IOException;
import java.util.Random;



public class PrintPosteriors {
	AbstractModel mhmm, mhmmb;

	AlignmentsSet forward, backward;

	public static void main(String[] args) throws IOException {
		PrintPosteriors pp = new PrintPosteriors();
		String corpusDescription = args[0];
		int size = Integer.parseInt(args[1]);
		int maxSentenceLen = Integer.parseInt(args[2]);
		String modelDir = args[3];
		String modelName = args[4];
		String decodingType = args[5];
		System.out.println("Corpus " + corpusDescription);
		System.out.println("Size " + size);
		System.out.println("Max Sentence Len " + maxSentenceLen);
		System.out.println("ModelDir " + modelDir);
		System.out.println("Model Name " + modelName);
		System.out.println("Decoding type " + decodingType);
		System.out.println("---");
		BilingualCorpus forwardCorpus = BilingualCorpus
				.getCorpusFromFileDescription(corpusDescription, size,
						maxSentenceLen);
		BilingualCorpus backwardCorpus = forwardCorpus.reverse();
		modelDir = modelDir + "/" + modelName + "/" + forwardCorpus.getName()
				+ "/" + size + "/model/";
		System.out.println("Loading Model");
		if (modelName.equalsIgnoreCase("baseline")) {
			pp.mhmm = HMM.loadModel(forwardCorpus, modelDir + "/MHMM/forward");
			pp.mhmmb = HMM.loadModel(backwardCorpus, modelDir
					+ "/MHMM/backward");
		} else if (modelName.equalsIgnoreCase("agreement")) {
			pp.mhmm = AgreementHMM.loadModel(forwardCorpus, backwardCorpus,
					modelDir + "/MHMM/");
			pp.mhmmb = ((AgreementHMM) pp.mhmm).getBackwardModel();
		} else {
			System.out.println("Unknown Model: " + modelName);
			System.exit(1);
		}
		System.out.println("Finished Loading Model");

		if (decodingType.equalsIgnoreCase("viter")) {
			pp.forward = pp.mhmm.viterbiAlignments(BilingualCorpus.TEST_CORPUS);
			pp.backward = pp.mhmmb
					.viterbiAlignments(BilingualCorpus.TEST_CORPUS);
		} else if (decodingType.equalsIgnoreCase("post-aer")) {
			float tresh = pp.mhmm.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,false);
			float tresh2 = pp.mhmmb
					.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,false);
			pp.forward = pp.mhmm.posteriorAlignments(
					BilingualCorpus.TEST_CORPUS, tresh,false,false);
			pp.backward = pp.mhmmb.posteriorAlignments(
					BilingualCorpus.TEST_CORPUS, tresh2,false,false);
		} else if (decodingType.startsWith("post-pts")) {
			// sample ~200 sentences from the training corpus
			System.out.println("Forward" + pp.mhmm);
			System.out.println("Backward" + pp.mhmmb);
			Random r = new Random(-9119629771948792409L);
			int totNumSent = forwardCorpus.getNumberOfTrainingSentences();
			int[] sample = new int[200];
			for (int i = 0; i < sample.length; i++) {
				sample[i] = r.nextInt(totNumSent);
			}
			// count number of aligned points with viterbi
			int numAlignmentsF = pp.mhmm.viterbiAlignments(sample,
					BilingualCorpus.TRAIN_CORPUS).numberOfAlignedPoints();
			int numAlignmentsB = pp.mhmmb.viterbiAlignments(sample,
					BilingualCorpus.TRAIN_CORPUS).numberOfAlignedPoints();
			// tune posterior to get that.
			float tresh = pp.mhmm.tuneThresholdNumPoints(numAlignmentsF,
					sample, BilingualCorpus.TRAIN_CORPUS,false);
			float tresh2 = pp.mhmmb.tuneThresholdNumPoints(numAlignmentsB,
					sample, BilingualCorpus.TRAIN_CORPUS,false);
			pp.forward = pp.mhmm.posteriorAlignments(
					BilingualCorpus.TEST_CORPUS, tresh,false,false);
			pp.backward = pp.mhmmb.posteriorAlignments(
					BilingualCorpus.TEST_CORPUS, tresh2,false,false);

		} else {
			System.out.println("Unknowon decoding");
		}

		System.out.println("\\documentclass{article} ");
		System.out.println("\\begin{document} ");
		System.out.println("\\section{Forward Alignments} ");
		System.out.println();
		System.out.println();
		AlignerOutputLatex.outputWithGoldPostLatex(pp.forward, forwardCorpus
				.getGold(), forwardCorpus, System.out, 10, "");
		System.out.println("Backward Alignments ");
		System.out.println();
		System.out.println();
		AlignerOutputLatex.outputWithGoldPostLatex(pp.backward, backwardCorpus
				.getGold(), backwardCorpus, System.out, 10, "");
		System.out.println();
		System.out.println();
		System.out.println("\\end{document} ");

	}
}
